# here put the import lib
import math
import sys

from PyQt5 import QtGui, QtWidgets

from general import plot_one_box, infer_img
from demo1 import Ui_MainWindow

import numpy as np
import cv2
# import time
# from random import uniform
from PyQt5.Qt import *
import onnxruntime as ort


class Open_Camera(QtWidgets.QMainWindow, Ui_MainWindow):
    def __init__(self):
        super(Open_Camera, self).__init__()
        self.setupUi(self)  # 创建窗体对象
        self.setWindowTitle('yolov5目标检测demo')
        self.init()

        self.openfile_name_model = 'yolov5n.onnx'  # 模型名称

        # self.so = ort.SessionOptions()  # 树莓派上保留以下两段代码,注释下面那行代码
        # self.net = ort.InferenceSession(self.openfile_name_model, self.so)

        # ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
        self.net = ort.InferenceSession(self.openfile_name_model,
                                        providers=['CUDAExecutionProvider'])  # 在树莓派上这里不需指定推理设备

        # 标签字典
        self.dic_labels = {0: 'person', 1: 'bicycle'}

        # 模型参数
        self.model_h = 320
        self.model_w = 320

    def init(self):
        # self.label.setScaledContents(True)  # 图片自适应
        # 定时器让其定时读取显示图片
        self.camera_timer = QTimer()
        self.camera_timer.timeout.connect(self.show_image)
        # 打开摄像头
        self.pushButton.clicked.connect(self.open_camera)
        # 关闭摄像头
        self.pushButton_2.clicked.connect(self.close_camera)

    '''开启摄像头'''
    def open_camera(self):
        self.cap = cv2.VideoCapture(0)  # 初始化摄像头
        self.camera_timer.start(40)  # 每40毫秒读取一次，即刷新率为25帧
        self.show_image()

    '''关闭摄像头'''
    def close_camera(self):
        self.cap.release()  # 释放摄像头
        self.label.clear()  # 清除label组件上的图片
        self.camera_timer.stop()  # 停止读取

    '''显示图片'''
    def show_image(self):
        flag, self.image = self.cap.read()  # 从视频流中读取图片

        if flag:
            # image_show = cv2.resize(self.image, (1280, 720))  # 把读到的帧的大小重新设置为 600*360
            image_show = self.image
            width, height = image_show.shape[:2]  # 行:宽，列:高
            # image_show = cv2.cvtColor(image_show, cv2.COLOR_BGR2RGB)  # opencv读的通道是BGR,要转成RGB
            # image_show = cv2.flip(image_show, 1)  # 水平翻转，因为摄像头拍的是镜像的。 用相机不用翻转
            # start 图片检测
            det_boxes, scores, ids = infer_img(image_show, self.net, self.model_h, self.model_w,
                                               thred_nms=0.5, thred_cond=0.75)
            # image_show = self.image
            for box, score, id in zip(det_boxes, scores, ids):
                label = '%s:%.2f' % (self.dic_labels[id], score)
                plot_one_box(box.astype(np.int16), image_show, color=(255, 0, 0), label=label, line_thickness=None)

        image_show = cv2.cvtColor(image_show, cv2.COLOR_BGR2RGB)  # opencv读的通道是BGR,要转成RGB
        # 把读取到的视频数据变成QImage形式(图片数据、高、宽、RGB颜色空间，三个通道各有2**8=256种颜色)
        self.showImage = QtGui.QImage(image_show.data, height, width, QImage.Format_RGB888)
        self.label.setPixmap(QPixmap.fromImage(self.showImage))  # 往显示视频的Label里显示QImage
        self.label.setScaledContents(True)  # 图片自适应


if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    ui = Open_Camera()
    ui.show()
    sys.exit(app.exec_())
